import os
from typing import List, Literal

from langchain_community.chat_models.zhipuai import ChatZhipuAI
from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, RemoveMessage
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph, START, END, MessagesState
from langgraph.prebuilt import ToolNode
from typing_extensions import TypedDict

os.environ["ZHIPUAI_API_KEY"] = "97738d4998b8732d707daf91a2b1c56d.2y6VKEuOlidwHDpI"

memory = MemorySaver()


class State(MessagesState):
    summary: str
    intention: str
    is_summary: bool


# 查询保单
@tool
def insurance_order_search(name: str) -> str:
    """可以根据用户的姓名和身份证号后4位查询用户在途的保单，当用户提供了上述信息之后，必须调用该方法获取保单信息"""
    return '''[
    {
        "长短险类型": "短险",
        "保险产品名称": "水滴接好孕·母婴保3.0（优享版）",
        "是否可水滴正常退": true,
        "投保人身份证后四位": "7241",
        "保单状态": 4,
        "正常退保预计退费金额": 238,
        "手机号": "13812345678",
        "姓名": "张三"
    },
    {
        "长短险类型": "短险",
        "保险产品名称": "水滴全能保癌无忧医疗险（焕新版）",
        "是否可水滴正常退": true,
        "投保人身份证后四位": "7241",
        "保单状态": 4,
        "正常退保预计退费金额": 238,
        "手机号": "13812345678",
        "姓名": "张三"
    },
    {
        "长短险类型": "短险",
        "保险产品名称": "水滴蓝海1号重疾险",
        "是否可水滴正常退": true,
        "投保人身份证后四位": "122X",
        "保单状态": 4,
        "正常退保预计退费金额": 96,
        "手机号": "13812345678",
        "姓名": "张三"
    },
    {
        "长短险类型": "短险",
        "保险产品名称": "水滴老年医疗险",
        "是否可水滴正常退": true,
        "投保人身份证后四位": "3527",
        "保单状态": 4,
        "正常退保预计退费金额": 132.39,
        "手机号": "13812345678",
        "姓名": "张三"
    }
]'''


# 保单退保
@tool
def insurance_order_cancel(policy_no: str) -> str:
    """根据用户提供的保单号进行退保操作"""
    return "已完成退保"


tools = [insurance_order_search, insurance_order_cancel]
tool_node = ToolNode(tools)
model = ChatZhipuAI(
    model='glm-4',
    temperature=0
).bind_tools(tools)


def get_context(messages: List[BaseMessage]) -> str:
    if len(messages) <= 0:
        return ""
    return "\n\n".join([message.type + ":" + message.content for message in messages])


def call_model(state: State):
    summary = state.get("summary", "")
    system_prompt = """你是一个保险客户专员，现在已经进入到了客户退保的流程，请你按照如下步骤引导用户操作。必须严格按照下面的步骤，一步一步的引导，不要提前暴露后续阶段的其他信息，也不能跨步骤进行。
    1、在退保的最开始阶段，要生成委婉的话术询问用户退保原因。
    2、在用户回复了退保原因之后，要求用委婉的话术确认用户是否确认退保，包含挽留之意，并且说明保险对用户的好处。
    3、如果用户确认要退保的话，需要优先询问用户，收集用户的姓名和身份证号后4位。如果用户确认不退保，则询问用户是否还有其他问题，如果没有其他问题的话，生成结束语给用户，并完结对话。
    4、根据用户提供的姓名和身份证号后4位，必须使用我提供的insurance_order_search工具查询用户目前在途的保单信息，并将工具返回的保单信息分条带数字展示，只展示保险产品名称即可。
    5、用户确认完退保的保单号之后，使用提供的insurance_order_cancel工具执行用户退保的操作。
    6、执行完用户的退保操作之后，询问用户是否还有其他问题，如果没有其他问题的话，生成结束语给用户，并完结对话。
    """
    if summary:
        messages = [SystemMessage(content=system_prompt + f"\n\n下面是历史上下文会话的总结：\n{summary}")] + state[
            "messages"]
    else:
        messages = [SystemMessage(content=system_prompt)] + state["messages"]
    response = model.invoke(messages)
    return {"messages": response}


def should_continue(state: State, config: RunnableConfig):
    last_message = state["messages"][-1]
    if not last_message.tool_calls:
        return "end"
    else:
        return "tools"


def should_summarize(state: State):
    messages = state["messages"]
    if len(messages) > 6:
        return "summarize_conversation"
    elif state["messages"][-1].tool_calls:
        return "tools"
    return "end"


def summarize_conversation(state: State):
    summary = state.get("summary", "")
    if summary:
        summary_message = (
            f"下面是之前总结的信息: {summary}\n\n"
            "根据上面总结的信息，以及新增的消息内容，再次进行总结:"
        )
    else:
        summary_message = "把上面的消息进行总结:"
    messages = state["messages"][:-2] + [HumanMessage(content=summary_message)]
    response = model.invoke(messages)
    return {"summary": response.content, "messages": [RemoveMessage(id=m.id) for m in state["messages"][:-2]],
            "is_summary": True}


sub_work_flow = StateGraph(State)
sub_work_flow.add_node("summarize_conversation", summarize_conversation)
sub_work_flow.add_node("call_model", call_model)
sub_work_flow.add_node("tool_node", tool_node)

sub_work_flow.add_edge(START, "call_model")
sub_work_flow.add_conditional_edges("call_model",
                                    should_summarize,
                                    {"summarize_conversation": "summarize_conversation", "tools": "tool_node",
                                     "end": END})
sub_work_flow.add_conditional_edges("summarize_conversation",
                                    should_continue,
                                    {"end": END, "tools": "tool_node"})
sub_work_flow.add_edge("tool_node", "call_model")

sub_graph = sub_work_flow.compile(checkpointer=memory)


class intention(TypedDict):
    intention: Literal["退保", "非退保其他保险业务", "结束会话"]


intention_recognition_prompt = """你是一个保险客户专员，能够结合用户的历史消息，识别出最新一条用户消息的核心意图。
历史消息总结:
{summary}

消息上下文
{context}

用户最新的消息
{user_message}
"""


def intention_recognition(state: State):
    prompt = format_prompt(state, intention_recognition_prompt)
    response = model.with_structured_output(intention).invoke(prompt)
    return {"intention": "其他" if response is None else response["intention"]}


def intention_condition(state: State):
    if state["intention"] == '退保':
        return "surrender"
    elif state["intention"] == '非退保其他保险业务':
        return "non_surrender"
    elif state["intention"] == "其他":
        return "salutation"
    else:
        return "end_session"


def format_prompt(state: State, prompt_template: str) -> str:
    context_messages = []
    if len(state["messages"]) > 1:
        context_messages = state["messages"][0:len(state["messages"]) - 1]
    context = ""
    for item in context_messages:
        context = context + item.type + ":" + item.content + "\n"
    prompt = prompt_template.format(
        summary=state.get("summary", "暂无"),
        context=context,
        user_message=state["messages"][len(state["messages"]) - 1].type + ":" + state["messages"][
            len(state["messages"]) - 1].content)
    return prompt


def non_surrender(state: State):
    return {"messages": [("ai", "好的，我帮你转人工客服进行处理")]}


def salutation(state: State):
    return {"messages": [("ai", "你好，我是水滴保险客服小依，有什么可以帮助到您的")]}


def end_session(state: State):
    return {"messages": [("ai", "好的，那我就先完结此次会话了，祝您生活愉快")]}


def should_delete_messages(state: State):
    if state.get("is_summary", False):
        return {"messages": [RemoveMessage(id=m.id) for m in state["messages"][:-2]], "is_summary": False}
    return {"is_summary": False}


work_flow = StateGraph(State)
work_flow.add_node("intention_recognition", intention_recognition)
work_flow.add_node("agent_executor", sub_graph)
work_flow.add_node("non_surrender", non_surrender)
work_flow.add_node("end_session", end_session)
work_flow.add_node("salutation", salutation)
work_flow.add_node("should_delete_messages", should_delete_messages)

work_flow.add_edge(START, "intention_recognition")
work_flow.add_conditional_edges("intention_recognition",
                                intention_condition,
                                {"surrender": "agent_executor",
                                 "non_surrender": "non_surrender",
                                 "salutation": "salutation",
                                 "end_session": "end_session"})
work_flow.add_edge("non_surrender", END)
work_flow.add_edge("end_session", END)
work_flow.add_edge("salutation", END)
work_flow.add_edge("agent_executor", "should_delete_messages")
work_flow.add_edge("should_delete_messages", END)

graph = work_flow.compile(checkpointer=memory)


def print_stream(stream):
    for s in stream:
        message = s["messages"][-1]
        if isinstance(message, tuple):
            print(message)
        else:
            message.pretty_print()

# config = {"configurable": {"thread_id": "1"}}
# print_stream(graph.stream({"messages": [("user", "小助手你好")]}, config, stream_mode="values"))
# print_stream(graph.stream({"messages": [("user", "退保")]}, config, stream_mode="values"))
# print_stream(graph.stream({"messages": [("user", "没啥用处了")]}, config, stream_mode="values"))
# print_stream(graph.stream({"messages": [("user", "确认退保")]}, config, stream_mode="values"))
# print_stream(graph.stream({"messages": [("user", "王亮，1234")]}, config, stream_mode="values"))
